package com.katze.common.ws.core;

import com.katze.common.ws.protocol.Message;
import com.katze.common.ws.protocol.MessageType;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.util.AttributeKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.Serializable;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

public class UserChannelGroup {
    private static final Logger log = LoggerFactory.getLogger(UserChannelGroup.class);

    public final static AttributeKey<Serializable> USER = AttributeKey.valueOf("user");
    private final static AttributeKey<MessageType> TYPE = AttributeKey.valueOf("type");
    private static final Map<Serializable, Map<MessageType, Channel>> userChannels = new ConcurrentHashMap<>(1024);
    private final ChannelFutureListener remover;

    public UserChannelGroup() {
        this.remover = future -> {
            Channel channel = future.channel();
            MessageType type = channel.attr(TYPE).get();
            Serializable user = channel.attr(USER).get();
            UserChannelGroup.this.remove(user, type);
        };
    }

    public void bind(Serializable user, Channel channel) {
        bind(user, MessageType.MESSAGE, channel);
    }

    private void bind(Serializable user, MessageType type, Channel channel) {
        channel.attr(USER).set(user); // 用户身份
        channel.attr(TYPE).set(type); // 通道类型

        // 注意：虽然通道关闭了，添加新通道操作在前，触发remove方法在后，此时会出现一个用户有多个通道
        Map<MessageType, Channel> channels = userChannels.compute(user, (k, v) -> {
            if (v == null) {
                v = new ConcurrentHashMap<>(4);
            }
            // 关闭旧通道（同类型）
            v.merge(type, channel, (v1, v2) -> {
                if (v1.isActive()) {
                    v1.close();
                }
                return v2;
            });
            channel.closeFuture().addListener(this.remover);
            return v;
        });
        log.info("{}@{}连接已保存, 总缓存: {}, 用户缓存: {}", type, user, userChannels.size(), channels.size());
    }

    public boolean isBound(Channel channel) {
        return channel.hasAttr(USER);
    }

    /**
     * 移除通道
     * @param user 用户
     */
    public void remove(Serializable user) {
        remove(user, MessageType.MESSAGE);
    }

    private void remove(Serializable user, MessageType type) {
        Map<MessageType, Channel> channels = userChannels.computeIfPresent(user, (k, v) -> {
            Channel channel = v.remove(type);
            if (channel != null) {
                channel.closeFuture().removeListener(this.remover);
            }
            return v.isEmpty() ? null : v;
        });
        log.info("{}@{}连接已移除, 总缓存: {}, 用户缓存: {}", type, user, userChannels.size(), channels == null? 0:channels.size());
    }


    /**
     * 发送消息给指定用户
     * @param user 用户
     * @param data 消息
     */
    public Consumer<Consumer<ChannelFuture>> send(Serializable user, Object data) {
        ChannelFuture future;
        Channel channel = userChannels.get(user).get(MessageType.MESSAGE);
        if (channel != null && channel.isActive()) {
            Message message = data instanceof Message ? (Message) data : new Message().setData(data);
            future = channel.writeAndFlush(message);
        } else {
            future = null;
            if (channel != null) {
                this.remove(user, MessageType.MESSAGE);
            }
        }
        return cs -> cs.accept(future);
    }

    private Consumer<Consumer<ChannelFuture>> send(Serializable user, MessageType type, Object data) {
        ChannelFuture future;
        Map<MessageType, Channel> channels = userChannels.get(user);
        Channel channel = channels != null? channels.get(type):null;
        if (channel != null && channel.isActive()) {
            Message message = data instanceof Message ? (Message) data : new Message().setData(data);
            future = channel.writeAndFlush(message);
        } else {
            future = null;
            if (channel != null) {
                this.remove(user, type);
            }
        }
        return cs -> cs.accept(future);
    }

    /**
     * 广播消息给所有用户
     * @param data 消息
     */
    public Consumer<Consumer<ChannelFuture>> broadcast(Object data) {
        Message message = data instanceof Message ? (Message) data : new Message().setData(data);
        return call -> {
            for (Map.Entry<Serializable, Map<MessageType, Channel>> entry : userChannels.entrySet()) {
                try {
                    Channel channel = entry.getValue().get(MessageType.MESSAGE);
                    if (!channel.isActive()) return;
                    ChannelFuture future = channel.writeAndFlush(message);
                    if (call != null) {
                        call.accept(future);
                    }
                } catch (Exception e) {
                    log.error("{}用户广播失败", entry.getKey(), e);
                }
            }
        };
    }

    private Consumer<Consumer<ChannelFuture>> broadcast(Object data, MessageType... types) {
        Message message = data instanceof Message ? (Message) data : new Message().setData(data);
        return call -> {
            for (Map.Entry<Serializable, Map<MessageType, Channel>> entry : userChannels.entrySet()) {
                try {
                    Map<MessageType, Channel> channels = entry.getValue();
                    for (MessageType type : types) {
                        Channel channel = channels.get(type);
                        if (!channel.isActive()) return;
                        ChannelFuture future = channel.writeAndFlush(message);
                        if (call != null) {
                            call.accept(future);
                        }
                    }
                } catch (Exception e) {
                    log.error("{}用户广播失败", entry.getKey(), e);
                }
            }
        };
    }
}
